package ace.oauth2.provider.common.support;

import lombok.SneakyThrows;
import org.redisson.api.*;
import org.springframework.security.oauth2.common.ExpiringOAuth2RefreshToken;
import org.springframework.security.oauth2.common.OAuth2AccessToken;
import org.springframework.security.oauth2.common.OAuth2RefreshToken;
import org.springframework.security.oauth2.provider.OAuth2Authentication;
import org.springframework.security.oauth2.provider.token.AuthenticationKeyGenerator;
import org.springframework.security.oauth2.provider.token.DefaultAuthenticationKeyGenerator;
import org.springframework.security.oauth2.provider.token.TokenStore;
import org.springframework.security.oauth2.provider.token.store.redis.RedisTokenStore;
import org.springframework.util.StringUtils;

import java.util.*;
import java.util.concurrent.TimeUnit;

/**
 * @author Caspar
 * @contract 279397942@qq.com
 * @create 2020/10/23
 * @description 功能与 {@link RedisTokenStore}一样，只是替换底层访问redis的组件{@link RedissonClient}
 */
public class RedissonTokenStore implements TokenStore {

    private static final String ACCESS = "access:";
    private static final String AUTH_TO_ACCESS = "auth_to_access:";
    private static final String AUTH = "auth:";
    private static final String REFRESH_AUTH = "refresh_auth:";
    private static final String ACCESS_TO_REFRESH = "access_to_refresh:";
    private static final String REFRESH = "refresh:";
    private static final String REFRESH_TO_ACCESS = "refresh_to_access:";
    private static final String CLIENT_ID_TO_ACCESS = "client_id_to_access:";
    private static final String UNAME_TO_ACCESS = "uname_to_access:";


    private AuthenticationKeyGenerator authenticationKeyGenerator = new DefaultAuthenticationKeyGenerator();

    private RedissonClient redissonClient;

    private String prefix = "";

    public RedissonTokenStore(RedissonClient redissonClient) {
        this.redissonClient = redissonClient;
    }

    public void setAuthenticationKeyGenerator(AuthenticationKeyGenerator authenticationKeyGenerator) {
        this.authenticationKeyGenerator = authenticationKeyGenerator;
    }


    public void setPrefix(String prefix) {
        this.prefix = prefix;
    }

    @Override
    public OAuth2AccessToken getAccessToken(OAuth2Authentication authentication) {
        String key = authenticationKeyGenerator.extractKey(authentication);
        String fullKey = buildKey(AUTH_TO_ACCESS + key);
        RBucket<OAuth2AccessToken> accessTokenRBucket = redissonClient.getBucket(fullKey);
        OAuth2AccessToken accessToken = accessTokenRBucket.get();
        if (Objects.isNull(accessToken)) {
            return null;
        }

        OAuth2Authentication storedAuthentication = readAuthentication(accessToken.getValue());
        if ((storedAuthentication == null || !key.equals(authenticationKeyGenerator.extractKey(storedAuthentication)))) {
            // Keep the stores consistent (maybe the same user is
            // represented by this authentication but the details have
            // changed)
            storeAccessToken(accessToken, authentication);
        }
        return accessToken;
    }

    @Override
    public OAuth2Authentication readAuthentication(OAuth2AccessToken token) {
        return readAuthentication(token.getValue());
    }

    @Override
    public OAuth2Authentication readAuthentication(String token) {
        String key = buildKey(AUTH + token);
        OAuth2Authentication auth = redissonClient.<OAuth2Authentication>getBucket(key).get();
        return auth;
    }

    @Override
    public OAuth2Authentication readAuthenticationForRefreshToken(OAuth2RefreshToken token) {
        return readAuthenticationForRefreshToken(token.getValue());
    }

    public OAuth2Authentication readAuthenticationForRefreshToken(String token) {
        String key = buildKey(REFRESH_AUTH + token);
        OAuth2Authentication auth = redissonClient.<OAuth2Authentication>getBucket(key).get();
        return auth;
    }

    @Override
    public void storeAccessToken(OAuth2AccessToken token, OAuth2Authentication authentication) {
        String accessKey = buildKey(ACCESS + token.getValue());
        String authKey = buildKey(AUTH + token.getValue());
        String authToAccessKey = buildKey(AUTH_TO_ACCESS + authenticationKeyGenerator.extractKey(authentication));
        String approvalKey = buildKey(UNAME_TO_ACCESS + getApprovalKey(authentication));
        String clientId = buildKey(CLIENT_ID_TO_ACCESS + authentication.getOAuth2Request().getClientId());

        RBatch rBatch = redissonClient.createBatch();

        RBucketAsync<OAuth2AccessToken> accessKeyWithTokenBucketAsync = rBatch.getBucket(accessKey);
        RBucketAsync<OAuth2AccessToken> authToAccessKeyWithTokenBucketAsync = rBatch.getBucket(authToAccessKey);
        RBucketAsync<OAuth2Authentication> authKeyWithAuthenticationBucketAsync = rBatch.getBucket(authKey);
        RSetAsync<OAuth2AccessToken> approvalKeyWithTokenSetAsync = rBatch.getSet(approvalKey);
        RSetAsync<OAuth2AccessToken> clientIdWithTokenSetAsync = rBatch.getSet(clientId);

        accessKeyWithTokenBucketAsync.setAsync(token);
        authToAccessKeyWithTokenBucketAsync.setAsync(token);
        authKeyWithAuthenticationBucketAsync.setAsync(authentication);
        if (!authentication.isClientOnly()) {
            approvalKeyWithTokenSetAsync.addAsync(token);
        }
        clientIdWithTokenSetAsync.addAsync(token);
        if (Objects.nonNull(token.getExpiration())) {
            accessKeyWithTokenBucketAsync.expireAsync(token.getExpiresIn(), TimeUnit.SECONDS);
            authToAccessKeyWithTokenBucketAsync.expireAsync(token.getExpiresIn(), TimeUnit.SECONDS);
            authKeyWithAuthenticationBucketAsync.expireAsync(token.getExpiresIn(), TimeUnit.SECONDS);
            approvalKeyWithTokenSetAsync.expireAsync(token.getExpiresIn(), TimeUnit.SECONDS);
            clientIdWithTokenSetAsync.expireAsync(token.getExpiresIn(), TimeUnit.SECONDS);
        }
        if (token.getRefreshToken() != null && token.getRefreshToken().getValue() != null) {
            String refreshToAccessKeyWithTokenValueKey = buildKey(REFRESH_TO_ACCESS + token.getRefreshToken().getValue());
            String accessToRefreshKeyWithRefreshTokenValue = buildKey(ACCESS_TO_REFRESH + token.getValue());

            RBucketAsync<String> refreshToAccessKeyWithTokenValueBucketAsync = rBatch.getBucket(refreshToAccessKeyWithTokenValueKey);
            RBucketAsync<String> accessToRefreshKeyWithRefreshTokenValueBucketAsync = rBatch.getBucket(accessToRefreshKeyWithRefreshTokenValue);

            refreshToAccessKeyWithTokenValueBucketAsync.setAsync(token.getValue());
            accessToRefreshKeyWithRefreshTokenValueBucketAsync.setAsync(token.getRefreshToken().getValue());

            if (token.getRefreshToken() instanceof ExpiringOAuth2RefreshToken) {
                ExpiringOAuth2RefreshToken expiringRefreshToken = (ExpiringOAuth2RefreshToken) token.getRefreshToken();
                Date expiration = expiringRefreshToken.getExpiration();
                if (expiration != null) {
                    int seconds = Long.valueOf((expiration.getTime() - System.currentTimeMillis()) / 1000L)
                            .intValue();
                    refreshToAccessKeyWithTokenValueBucketAsync.expireAsync(seconds, TimeUnit.SECONDS);
                    accessToRefreshKeyWithRefreshTokenValueBucketAsync.expireAsync(seconds, TimeUnit.SECONDS);
                }
            }
        }
        rBatch.execute();
    }

    private static String getApprovalKey(OAuth2Authentication authentication) {
        String userName = authentication.getUserAuthentication() == null ? ""
                : authentication.getUserAuthentication().getName();
        return getApprovalKey(authentication.getOAuth2Request().getClientId(), userName);
    }

    private static String getApprovalKey(String clientId, String userName) {
        return clientId + (userName == null ? "" : ":" + userName);
    }

    @Override
    public void removeAccessToken(OAuth2AccessToken accessToken) {
        removeAccessToken(accessToken.getValue());
    }

    @Override
    public OAuth2AccessToken readAccessToken(String tokenValue) {
        String key = buildKey(ACCESS + tokenValue);
        RBucket<OAuth2AccessToken> rBucket = redissonClient.getBucket(key);
        return rBucket.get();
    }

    @SneakyThrows
    public void removeAccessToken(String accessTokenValue) {
        String accessKey = buildKey(ACCESS + accessTokenValue);
        String authKey = buildKey(AUTH + accessTokenValue);
        String accessToRefreshKey = buildKey(ACCESS_TO_REFRESH + accessTokenValue);

        RBatch rBatch = redissonClient.createBatch();
        RBucketAsync<OAuth2AccessToken> accessKeyWithAccessTokenBucketAsync = rBatch.getBucket(accessKey);
        RBucketAsync<OAuth2Authentication> authKeyWithAuthBucketAsync = rBatch.getBucket(authKey);
        RBucketAsync<String> accessToRefreshKeyBucketAsync = rBatch.getBucket(accessToRefreshKey);

        RFuture<OAuth2AccessToken> oAuth2AccessTokenRFuture = accessKeyWithAccessTokenBucketAsync.getAsync();
        RFuture<OAuth2Authentication> oAuth2AuthenticationRFuture = authKeyWithAuthBucketAsync.getAsync();

        accessKeyWithAccessTokenBucketAsync.deleteAsync();
        authKeyWithAuthBucketAsync.deleteAsync();
        accessToRefreshKeyBucketAsync.deleteAsync();

        rBatch.execute();

        OAuth2AccessToken accessToken = oAuth2AccessTokenRFuture.get();
        OAuth2Authentication authentication = oAuth2AuthenticationRFuture.get();
        if (authentication != null) {
            String key = authenticationKeyGenerator.extractKey(authentication);
            String authToAccessKey = buildKey(AUTH_TO_ACCESS + key);
            String unameKey = buildKey(UNAME_TO_ACCESS + getApprovalKey(authentication));
            String clientId = buildKey(CLIENT_ID_TO_ACCESS + authentication.getOAuth2Request().getClientId());

            rBatch = redissonClient.createBatch();
            rBatch.getBucket(authToAccessKey).deleteAsync();
            rBatch.getSet(unameKey).removeAsync(accessToken);
            rBatch.getSet(clientId).removeAsync(accessToken);
            rBatch.getBucket(ACCESS + key).deleteAsync();
            rBatch.execute();
        }
    }

    @Override
    public void storeRefreshToken(OAuth2RefreshToken refreshToken, OAuth2Authentication authentication) {
        String refreshKey = buildKey(REFRESH + refreshToken.getValue());
        String refreshAuthKey = buildKey(REFRESH_AUTH + refreshToken.getValue());

        RBatch rBatch = redissonClient.createBatch();
        RBucketAsync<OAuth2RefreshToken> refreshTokenRBucketAsync = rBatch.getBucket(refreshKey);
        RBucketAsync<OAuth2Authentication> authenticationRBucketAsync = rBatch.getBucket(refreshAuthKey);

        refreshTokenRBucketAsync.setAsync(refreshToken);
        authenticationRBucketAsync.setAsync(authentication);
        if (refreshToken instanceof ExpiringOAuth2RefreshToken) {
            ExpiringOAuth2RefreshToken expiringRefreshToken = (ExpiringOAuth2RefreshToken) refreshToken;
            Date expiration = expiringRefreshToken.getExpiration();
            if (expiration != null) {
                int seconds = Long.valueOf((expiration.getTime() - System.currentTimeMillis()) / 1000L)
                        .intValue();
                refreshTokenRBucketAsync.expireAsync(seconds, TimeUnit.SECONDS);
                authenticationRBucketAsync.expireAsync(seconds, TimeUnit.SECONDS);
            }
        }
        rBatch.execute();
    }

    @Override
    public OAuth2RefreshToken readRefreshToken(String tokenValue) {
        String key = buildKey(REFRESH + tokenValue);
        OAuth2RefreshToken refreshToken = redissonClient.<OAuth2RefreshToken>getBucket(key).get();
        return refreshToken;
    }

    @Override
    public void removeRefreshToken(OAuth2RefreshToken refreshToken) {
        removeRefreshToken(refreshToken.getValue());
    }

    public void removeRefreshToken(String tokenValue) {
        String refreshKey = buildKey(REFRESH + tokenValue);
        String refreshAuthKey = buildKey(REFRESH_AUTH + tokenValue);
        String refresh2AccessKey = buildKey(REFRESH_TO_ACCESS + tokenValue);
        String access2RefreshKey = buildKey(ACCESS_TO_REFRESH + tokenValue);

        RBatch rBatch = redissonClient.createBatch();
        rBatch.getBucket(refreshKey).deleteAsync();
        rBatch.getBucket(refreshAuthKey).deleteAsync();
        rBatch.getBucket(refresh2AccessKey).deleteAsync();
        rBatch.getBucket(access2RefreshKey).deleteAsync();
        rBatch.execute();
    }

    @Override
    public void removeAccessTokenUsingRefreshToken(OAuth2RefreshToken refreshToken) {
        removeAccessTokenUsingRefreshToken(refreshToken.getValue());
    }

    private void removeAccessTokenUsingRefreshToken(String refreshToken) {
        String key = buildKey(REFRESH_TO_ACCESS + refreshToken);
        String accessToken = redissonClient.<String>getBucket(key).getAndDelete();
        if (StringUtils.isEmpty(accessToken)) {
            return;
        }
        removeAccessToken(accessToken);
    }

    @Override
    public Collection<OAuth2AccessToken> findTokensByClientIdAndUserName(String clientId, String userName) {
        String approvalKey = buildKey(UNAME_TO_ACCESS + getApprovalKey(clientId, userName));
        RSet<OAuth2AccessToken> oAuth2AccessTokenRSet = redissonClient.getSet(approvalKey);
        Set<OAuth2AccessToken> accessTokenSet = oAuth2AccessTokenRSet.readAll();
        if (accessTokenSet == null || accessTokenSet.size() == 0) {
            return Collections.EMPTY_LIST;
        }
        return Collections.unmodifiableCollection(accessTokenSet);
    }

    @Override
    public Collection<OAuth2AccessToken> findTokensByClientId(String clientId) {
        String key = buildKey(CLIENT_ID_TO_ACCESS + clientId);
        RSet<OAuth2AccessToken> accessTokenRSet = redissonClient.getSet(key);
        Set<OAuth2AccessToken> accessTokenSet = accessTokenRSet.readAll();
        if (accessTokenSet == null || accessTokenSet.size() == 0) {
            return Collections.EMPTY_LIST;
        }
        return Collections.unmodifiableCollection(accessTokenSet);
    }

    private String buildKey(String key) {
        return prefix + key;
    }
}